#include <iostream>
#include <vector>
#include <Utils.h>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/io/pcd_io.h>
#include <pcl/kdtree/kdtree_flann.h>
#include <Eigen/Core>
#include <Eigen/Dense>
using namespace std;
namespace xio=lyxutils::io;
namespace xos=lyxutils::os;
namespace np=lyxutils::numpy;
namespace dbg=lyxutils::debug;

void printUsage(){
    cout<<"usage:cal_normal pc.txt|pc.pcd  =>calculate point cloud normals"<<endl;
    cout<<"      pc.txt=>point cloud file.first three columns will be used as x,y,z"<<endl;
    cout<<"      12 columns:lambda1,lambda2,lambda3,n1x,n1y,n1z,n2x,n2y,n2z,n3x,n3y,n3z"<<endl;
    cout<<"      will be appended to data table,where lambda1>=lambda2>=lambda3"<<endl;
    cout<<"      and (n{i}x,n{i}y,n{i}z) represents eigen vector corresponding to lambda{i}"<<endl;
    cout<<"      note=>if pcd file is used,the input point cloud type will be treated as PointXYZ,"<<endl;
    cout<<"      output point cloud type will be PointNormal(and preserve mode will be auto disabled)"<<endl;
    cout<<"argument:"<<endl;
    cout<<"  -h            print this help message"<<endl;
    cout<<"  -o=file.txt   set output file name as file.txt(default is output.txt|output.pcd)"<<endl;
    cout<<"  -r=radius     set neighbor search radius(default is 1)"<<endl;
    cout<<"  -k=num        set number of nereast neighbors as num(default is 10)"<<endl;
    cout<<"  -m=mode       set neighbor search mode(choices:k,r,auto-default is auto)"<<endl;
    cout<<"                if -m=k:use k nereast eighbor search;if -m=r:use radius search"<<endl;
    cout<<"                if -m=auto:use radius search unless #points in radius is less than -k=num"<<endl;
    cout<<"  -p            only first three columns:x,y,z are used but the input data may contain "<<endl;
    cout<<"                extra columns;this option determines whether to preserve the extra columns or not"<<endl;
}

using Vector3f=lyxutils::Vector3<float>;

void planeFit(const vector<vector<float> > &pts, Vector3f &sigma, vector<Vector3f> &eigVecs){
    vector<vector<float> > covariance;
    np::cov(pts,covariance);
    Eigen::Matrix3d cov_mat;
    for(int i=0;i<3;++i)for(int j=0;j<3;++j)cov_mat(i,j)=covariance[i][j];
    Eigen::EigenSolver<Eigen::Matrix3d> es(cov_mat);
    Eigen::Vector3d eigs=es.pseudoEigenvalueMatrix().diagonal();
    Eigen::Matrix3d vecs=es.pseudoEigenvectors();
    double minEig=eigs(0),maxEig=eigs(0);
    int maxIdx=0,minIdx=0,midIdx=0;
    for(int i=1;i<3;++i){
        if(minEig>eigs(i)){
            minIdx=i;
            minEig=eigs(i);
        }
        if(maxEig<eigs(i)){
            maxIdx=i;
            maxEig=eigs(i);
        }
    }
    if(minIdx==maxIdx){
        maxIdx=0;
        midIdx=1;
        minIdx=2;
    }
    else{
        midIdx=3-minIdx-maxIdx;
    }
    sigma={(float)eigs(maxIdx),(float)eigs(midIdx),(float)eigs(minIdx)};
    eigVecs.clear();
    eigVecs.push_back({float(vecs(0,maxIdx)),float(vecs(1,maxIdx)),float(vecs(2,maxIdx))});
    eigVecs.push_back({float(vecs(0,midIdx)),float(vecs(1,midIdx)),float(vecs(2,midIdx))});
    eigVecs.push_back({float(vecs(0,minIdx)),float(vecs(1,minIdx)),float(vecs(2,minIdx))});
}

int main(int argc, char**argv) {
    xio::CLParser parser;
    parser.parse(argc,argv);
    vector<string> paras=parser.getParameters();
    if(paras.size()!=1 || parser.hasOption("h")){
        printUsage();
        return 1;
    }
    //deal with cmd options
    bool isPcd=(xos::suffix(paras[0])=="pcd");
    float radius(1.0);
    int K(10);
    int errorLimit=4;
    string outputFile=isPcd?"output.pcd":"output.txt";
    string mode="auto";
    try{
        radius=parser.hasOption("r")?stof(parser.getOptionArg("r")):1.0;
    }catch(invalid_argument ia){
        cout<<"invalid argument:"<<ia.what()<<endl;
    }
    try{
        K=parser.hasOption("k")?stoi(parser.getOptionArg("k")):10;
        if(K<errorLimit)throw invalid_argument("K should be at least "+errorLimit);
    }catch(invalid_argument ia){
        cout<<"invalid argument:"<<ia.what()<<endl;
    }
    try{
        mode=parser.hasOption("m")?parser.getOptionArg("m"):"auto";
    }catch(invalid_argument ia){
        cout<<"invalid argument:"<<ia.what()<<endl;
    }
    try{
        outputFile=parser.hasOption("o")?parser.getOptionArg("o"):outputFile;
    }catch(invalid_argument ia){
        cout<<"invalid argument:"<<ia.what()<<endl;
    }
    bool preserve=parser.hasOption("p");
    //do real stuff
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZ>);
    pcl::PointCloud<pcl::PointNormal>::Ptr cloudNormal(new pcl::PointCloud<pcl::PointNormal>);
    vector<vector<float> > other_fields;
    size_t lines=dbg::estimate_file_rows(paras[0]);
    // xio::read_pts_field(paras[0],cloud->points,other_fields,",",cout);
    if(xos::suffix(paras[0])=="pcd" && preserve){
        cout<<"note:input pcd file will be treated as PointXYZ cloud,extra fields(if exist) won't be preserved!"<<endl;
        preserve=false;
    }
    cout<<"reading file..."<<endl;
    int count=1;
    if(preserve){
        xio::read_csv_customize(paras[0],[&](const vector<float> &row){
            cloud->points.push_back({row[0],row[1],row[2]});
            vector<float> tmp;
            for(int i=3;i<row.size();++i){
                tmp.push_back(row[i]);
            }
            other_fields.push_back(tmp);
            xio::consoleProgressBar(count++,lines);
        });
    }
    else{
        if(xos::suffix(paras[0])=="pcd"){
            pcl::io::loadPCDFile(xos::abspath(paras[0]),*cloud);
        }
        else if(xos::suffix(paras[0])=="txt"){
            xio::read_csv_customize(paras[0],[&](const vector<float> &row){
                cloud->points.push_back({row[0],row[1],row[2]});
                other_fields.push_back(vector<float>());
                xio::consoleProgressBar(count++,lines);
            });
        }
    }
    if(lines>cloud->points.size()){
        //over estimate file lines causing progress bar not finished
        xio::consoleProgressBar(100,100);
    }
    cout<<"read point cloud has "<<cloud->size();
    if(preserve)cout<<",other fields:"<<other_fields.size()<<" rows and "<<other_fields[0].size()<<" cols"<<endl;
    else cout<<endl;
    cout<<"creating tree..."<<endl;
    pcl::KdTreeFLANN<pcl::PointXYZ>::Ptr tree(new pcl::KdTreeFLANN<pcl::PointXYZ>);
    tree->setInputCloud(cloud);
    cout<<"calculating start at:"<<dbg::getTimeString()<<endl;
    for(int i=0;i<cloud->size();++i){
        vector<int> idx;
        vector<float> sqrDis;
        pcl::PointXYZ pt=cloud->points[i];
        if(mode=="r"||mode=="auto"){
            tree->radiusSearch(pt,radius,idx,sqrDis);
        }
        if(mode=="k"||(mode=="auto" && idx.size()<K)){
            tree->nearestKSearch(pt,K,idx,sqrDis);
        }
        if(idx.size()<errorLimit){
            cout<<"error:number of neighbors not enough(<"<<errorLimit<<"),the point is:("<<pt.x<<","<<pt.y<<","<<pt.z<<")"<<endl;
            return 1;
        }
        Vector3f sigma;
        vector<Vector3f> vec;
        vector<vector<float> > pts;
        for(int j=0;j<idx.size();++j){
            pcl::PointXYZ pt=cloud->points[idx[j]];
            pts.push_back(vector<float>());
            (pts.end()-1)->push_back(pt.x);
            (pts.end()-1)->push_back(pt.y);
            (pts.end()-1)->push_back(pt.z);
        }
        planeFit(pts,sigma,vec);
        if(isPcd){
            pcl::PointNormal pTemp;
            pTemp.x=pt.x;
            pTemp.y=pt.y;
            pTemp.z=pt.z;
            pTemp.normal_x=vec[2].x;
            pTemp.normal_y=vec[2].y;
            pTemp.normal_z=vec[2].z;
            cloudNormal->points.push_back(pTemp);
        }
        else{
            for(int j=0;j<3;++j)other_fields[i].push_back(sigma[j]);
            for(int j=0;j<3;++j)for(int k=0;k<3;++k)other_fields[i].push_back(vec[j][k]);
        }
        xio::consoleProgressBar(i+1,cloud->size());
    }
    if(isPcd){
        cloudNormal->height=1;
        cloudNormal->width=cloudNormal->points.size();
    }
    cout<<"calculating end at:"<<dbg::getTimeString()<<endl;
    cout<<"writing to file:"<<outputFile<<endl;
    if(isPcd){
        pcl::io::savePCDFile(outputFile,*cloudNormal);
    }
    else if(xos::suffix(paras[0])=="txt"){
        ofstream ofs(outputFile);
        ofs << std::setprecision(8);
        for(int i=0;i<cloud->size();++i){
            pcl::PointXYZ pt=cloud->points[i];
            ofs<<pt.x<<","<<pt.y<<","<<pt.z;
            for(int j=0;j<other_fields[i].size();++j){
                ofs<<","<<other_fields[i][j];
            }
            ofs<<endl;
        }
        ofs.close();
    }
    cout<<"succeed and exit."<<endl;
}
